/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.sql.impl.rel;



import com.bff.gaia.unified.sdk.extensions.sql.impl.transform.UnifiedSetOperatorsTransforms;

import com.bff.gaia.unified.sdk.transforms.MapElements;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.transforms.ParDo;

import com.bff.gaia.unified.sdk.transforms.join.CoGbkResult;

import com.bff.gaia.unified.sdk.transforms.join.CoGroupByKey;

import com.bff.gaia.unified.sdk.transforms.join.KeyedPCollectionTuple;

import com.bff.gaia.unified.sdk.transforms.windowing.WindowFn;

import com.bff.gaia.unified.sdk.values.*;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.PCollectionList;

import com.bff.gaia.unified.sdk.values.Row;

import com.bff.gaia.unified.sdk.values.TupleTag;



import java.io.Serializable;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;



/**

 * Delegate for Set operators: {@code UnifiedUnionRel}, {@code UnifiedIntersectRel} and {@code

 * UnifiedMinusRel}.

 */

public class UnifiedSetOperatorRelBase extends PTransform<PCollectionList<Row>, PCollection<Row>> {

  /** Set operator type. */

  public enum OpType implements Serializable {

    UNION,

    INTERSECT,

    MINUS

  }



  private UnifiedRelNode unifiedRelNode;

  private boolean all;

  private OpType opType;



  public UnifiedSetOperatorRelBase(UnifiedRelNode unifiedRelNode, OpType opType, boolean all) {

    this.unifiedRelNode = unifiedRelNode;

    this.opType = opType;

    this.all = all;

  }



  @Override

  public PCollection<Row> expand(PCollectionList<Row> inputs) {

    checkArgument(

        inputs.size() == 2,

        "Wrong number of arguments to %s: %s",

        unifiedRelNode.getClass().getSimpleName(),

        inputs);

    PCollection<Row> leftRows = inputs.get(0);

    PCollection<Row> rightRows = inputs.get(1);



    WindowFn leftWindow = leftRows.getWindowingStrategy().getWindowFn();

    WindowFn rightWindow = rightRows.getWindowingStrategy().getWindowFn();

    if (!leftWindow.isCompatible(rightWindow)) {

      throw new IllegalArgumentException(

          "inputs of "

              + opType

              + " have different window strategy: "

              + leftWindow

              + " VS "

              + rightWindow);

    }



    final TupleTag<Row> leftTag = new TupleTag<>();

    final TupleTag<Row> rightTag = new TupleTag<>();



    // co-group

    PCollection<KV<Row, CoGbkResult>> coGbkResultCollection =

        KeyedPCollectionTuple.of(

                leftTag,

                leftRows.apply(

                    "CreateLeftIndex",

                    MapElements.via(new UnifiedSetOperatorsTransforms.UnifiedSqlRow2KvFn())))

            .and(

                rightTag,

                rightRows.apply(

                    "CreateRightIndex",

                    MapElements.via(new UnifiedSetOperatorsTransforms.UnifiedSqlRow2KvFn())))

            .apply(CoGroupByKey.create());

    return coGbkResultCollection.apply(

        ParDo.of(

            new UnifiedSetOperatorsTransforms.SetOperatorFilteringDoFn(

                leftTag, rightTag, opType, all)));

  }

}